package cn.itcast.flink.base;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * Author itcast
 * Date 2021/7/27 15:58
 * Desc 将生成的 0 - 100 之间的数字过滤掉小于10的元素，均匀的分不到三个分区上，使用重分布和不适用重分布的区别
 */
public class RebalanceTran {
    public static void main(String[] args) throws Exception {
        //1.env 设置并行度为3
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        //分到三个分区中
        env.setParallelism(3);
        //2.source fromSequence 0-100
        DataStreamSource<Long> source = env.fromSequence(0, 100);
        //3.Transformation
        //下面的操作相当于将数据随机分配一下,有可能出现数据倾斜，过滤出来大于10
        SingleOutputStreamOperator<Long> filterDataStream = source.filter(n -> n > 10);
        //3.1 接下来使用map操作,将Long数据转为(分区编号/子任务编号, 数据)
        //重分区就是要统计每个分区，每个subtask，每个线程处理数据的个数相等，只需要统计出来每个subtask处理的数据的个数
        //1 -> 30  2->30  3->30
            SingleOutputStreamOperator<Tuple2<Integer, Long>> mapDataStream = filterDataStream.map(new RichMapFunction<Long, Tuple2<Integer/**CPU的线程ID*/, Long/*1*/>>() {
            @Override
            //通过getRuntimeContext获取到任务Index
            //返回Tuple2(任务Index,1)
            //按照子任务id/分区编号分组，统计每个子任务/分区中有几个元素
            public Tuple2<Integer, Long> map(Long value) throws Exception {
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                return Tuple2.of(indexOfThisSubtask, 1L);
            }
        });
        SingleOutputStreamOperator<Tuple2<Integer, Long>> result1 = mapDataStream.keyBy(t -> t.f0)
                .sum(1);
        //3.2 重新执行以上操作在filter之后先 rebalance 再map
        SingleOutputStreamOperator<Tuple2<Integer, Long>> rebalanceDataStream = filterDataStream
                .rebalance()
                .map(new RichMapFunction<Long, Tuple2<Integer/**CPU的线程ID*/, Long/*1*/>>() {
            @Override
            //通过getRuntimeContext获取到任务Index
            //返回Tuple2(任务Index,1)
            //按照子任务id/分区编号分组，统计每个子任务/分区中有几个元素
            public Tuple2<Integer, Long> map(Long value) throws Exception {
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                return Tuple2.of(indexOfThisSubtask, 1L);
            }
        });
        SingleOutputStreamOperator<Tuple2<Integer, Long>> result2 = rebalanceDataStream
                .keyBy(t -> t.f0)
                .sum(1);
        //4.sink
        //result1.print();//有可能出现数据倾斜
        result2.print();
        //5.execute
        env.execute();
    }
}
